#! /usr/bin/env python
#coding=utf-8

import os
import codecs

# https://www.jianshu.com/p/8203457a11cc

class VirtualMemroyArea(dict):
	VMA_SIZE_KEYS = ("Size", "Rss", "Pss", "Swap", "SwapPss", "Shared_Clean", "Shared_Dirty", "Private_Clean", "Private_Dirty", "Referenced", "Anonymous", "KernelPageSize", "MMUPageSize", "LazyFree", "AnonHugePages", "ShmemPmdMapped", "Shared_Hugetlb", "Private_Hugetlb", "Locked", "THPeligible", "FilePmdMapped" )

	VMA_ADDR_KEYS = ("start", "end", "perm", "offset", "dev", "idx")
	#7fab708000-7fab71c000 r--p 00000000 fd:00 6179  /system/lib64/libpluginmanager.z.so
	#start      end        perm offset   dev   idx   (file)
	def __init__(self, parts):
		for i, k in enumerate(VirtualMemroyArea.VMA_ADDR_KEYS):
			if i < 2: # Skip addr start and end
				continue
			self[k] = parts[i-1]
		addrs = parts[0].split("-")
		self["start"] = addrs[0]
		self["end"] = addrs[1]

		# Initialize sizes
		for k in VirtualMemroyArea.VMA_SIZE_KEYS:
			self[k] = 0

class SmapsObject(dict):
	def __init__(self, name):
		self["name"] = name
		self["vmas"] = 0
		self._vmas = []

	def add_vma(self, parts):
		vma = VirtualMemroyArea(parts)
		self._vmas.append(vma)

	def add_vma_info(self, key, val):
		#print("%s %s" % (key, val))
		vma = self._vmas[-1]
		vma[key] = val

	def do_summary(self):
		for k in VirtualMemroyArea.VMA_SIZE_KEYS:
			self[k] = 0

		for vma in self._vmas:
			for k in VirtualMemroyArea.VMA_SIZE_KEYS:
				self[k] = self[k] + vma[k]

		self["vmas"] = len(self._vmas)

	def get_all(self):
		return self._vmas

class ProcessSmaps(dict):
	def __init__(self, filename):
		self._objects = []
		self._obj_dict = {}
		self["category"] = "posix"
		self["objects"] = 0
		self["vmas"] = 0
		for k in VirtualMemroyArea.VMA_SIZE_KEYS:
			self[k] = 0
		self.__load(filename)

	def get_all(self):
		return self._objects

	def _parse_object_line(self, last_obj, parts):
		if len(parts) <= len(VirtualMemroyArea.VMA_ADDR_KEYS) - 1:
			last_obj.add_vma(parts)
			return last_obj
		name = " ".join(parts[len(VirtualMemroyArea.VMA_ADDR_KEYS)-1:])
		if name.startswith("/"):
			name = name[1:]
		if name == "system/bin/sa_main":
			self["category"] = "sa"
		elif name == "vendor/bin/hdf_devmgr":
			self["category"] = "hdi_service"
		elif name == "system/bin/appspawn":
			self["category"] = "app"
		if last_obj and last_obj["name"] == name:
			last_obj.add_vma(parts)
			return last_obj
		if name in self._obj_dict:
			obj = self._obj_dict[name]
		else:
			obj = SmapsObject(name)
			self._objects.append(obj)
			self._obj_dict[name] = obj
			self["objects"] = self["objects"] + 1
		obj.add_vma(parts)
		return obj

	def _parse_size_and_flags(self, last_obj, parts, filename):
		if not last_obj:
			return
		key = parts[0]
		if key.endswith(":"):
			key = key[:-1]
		if key in VirtualMemroyArea.VMA_SIZE_KEYS:
			last_obj.add_vma_info(key, int(parts[1]))
			return
		if key == "VmFlags":
			last_obj.add_vma_info("VmFlags", " ".join(parts[1:]))
		elif not key == "Name":
			print("[%s] %s: %s" %(filename, key, " ".join(parts[1:])))

	def __load(self, filename):
		with codecs.open(filename, 'r', 'utf-8') as f:
			lastObj = None
			for line in f.readlines():
				line = line.strip()
				if "shell cat" in line or "No such file or directory" in line or "Not support" in line:
					continue
				parts = line.split()
				key = parts[0]
				if key[-1] == ":" or key.find("-") < 0:
					if not lastObj:
						continue
					self._parse_size_and_flags(lastObj, parts, filename)
					continue

				# parse a new object
				lastObj = self._parse_object_line(lastObj, parts)

		for obj in self._objects:
			obj.do_summary()
			for k in VirtualMemroyArea.VMA_SIZE_KEYS:
				self[k] = self[k] + obj[k]
			self["vmas"] = self["vmas"] + obj["vmas"]

	def report(self, f):
		f.write("<items>\n")
		for idx, obj in enumerate(self._objects):
			f.write("<item>\n")
			f.write("<id>%d</id>\n" % (idx + 1))
			f.write("<name>%s</name>\n" % obj["name"])
			f.write("<vmas>%s</vmas>\n" % obj["vmas"])
			for k in VirtualMemroyArea.VMA_SIZE_KEYS:
				f.write("<%s>%s</%s>\n" % (k, obj[k], k))
			f.write("</item>\n")
		f.write("</items>\n")

if __name__ == "__main__":
	import sys
	import traceback

	if len(sys.argv) <= 1:
		print("smaps.py smaps_file1 smaps_file2 ...")
		sys.exit()

	for smap in sys.argv[1:]:
		try:
			print("Parsing smaps file: %s" % smap)
			p = ProcessSmaps(smap)
			out = os.path.splitext(smap)[0] + ".xml"
			with open(out, "w") as f:
				p.report(f)
			print("  parsed smaps file: %s" % out)
		except:
			traceback.print_exc()

